from flask import Flask, render_template
from flask_socketio import SocketIO, emit
import paramiko
import threading
import time

class TerminalServer:
    def __init__(self):
        self.app = Flask(__name__)
        self.app.config['SECRET_KEY'] = 'terminal-secret-key'
        self.socketio = SocketIO(self.app, cors_allowed_origins="*")
        
        # 存储SSH连接
        self.connections = {}
        
        self.setup_routes()
        self.setup_socketio()
    
    def setup_routes(self):
        @self.app.route('/')
        def index():
            return render_template('index.html')
    
    def setup_socketio(self):
        @self.socketio.on('connect')
        def handle_connect():
            print(f'Client connected: {request.sid}')
            emit('output', {'data': '=== Terminal Ready ===\n'})
        
        @self.socketio.on('disconnect')
        def handle_disconnect():
            sid = request.sid
            print(f'Client disconnected: {sid}')
            if sid in self.connections:
                try:
                    self.connections[sid]['ssh'].close()
                except:
                    pass
                del self.connections[sid]
        
        @self.socketio.on('command')
        def handle_command(data):
            from flask import request
            sid = request.sid
            cmd = data.get('cmd', '').strip()
            
            if cmd.startswith('connect '):
                self.handle_ssh_connect(sid, cmd)
            elif cmd == 'disconnect':
                self.handle_disconnect(sid)
            elif cmd == 'help':
                self.show_help(sid)
            else:
                self.execute_command(sid, cmd)
    
    def handle_ssh_connect(self, sid, cmd):
        try:
            # 解析: connect user@host password
            parts = cmd.split(' ', 2)
            if len(parts) < 3:
                emit('output', {'data': 'Usage: connect user@host password\n'})
                return
            
            user_host = parts[1]
            password = parts[2]
            
            if '@' not in user_host:
                emit('output', {'data': 'Format: user@host\n'})
                return
            
            user, host = user_host.split('@', 1)
            
            emit('output', {'data': f'Connecting to {user}@{host}...\n'})
            
            # 创建SSH连接
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(hostname=host, username=user, password=password, timeout=10)
            
            self.connections[sid] = {'ssh': ssh}
            emit('output', {'data': f'Connected to {user}@{host}\n'})
            
        except Exception as e:
            emit('output', {'data': f'Connection failed: {str(e)}\n'})
    
    def handle_disconnect(self, sid):
        if sid in self.connections:
            try:
                self.connections[sid]['ssh'].close()
                del self.connections[sid]
                emit('output', {'data': 'Disconnected\n'})
            except:
                emit('output', {'data': 'Disconnect failed\n'})
        else:
            emit('output', {'data': 'No active connection\n'})
    
    def execute_command(self, sid, cmd):
        if sid not in self.connections:
            emit('output', {'data': 'Not connected. Use: connect user@host password\n'})
            return
        
        try:
            ssh = self.connections[sid]['ssh']
            stdin, stdout, stderr = ssh.exec_command(cmd)
            
            output = stdout.read().decode('utf-8')
            error = stderr.read().decode('utf-8')
            
            if output:
                emit('output', {'data': output})
            if error:
                emit('output', {'data': f'Error: {error}'})
                
        except Exception as e:
            emit('output', {'data': f'Command failed: {str(e)}\n'})
    
    def show_help(self, sid):
        help_text = """Available commands:
connect user@host password  - Connect to SSH server
disconnect                  - Disconnect from server
help                       - Show this help
Any Linux command after connecting

Example: connect root@192.168.1.100 mypassword
"""
        emit('output', {'data': help_text})
    
    def run(self, host='127.0.0.1', port=5000, debug=True):
        self.socketio.run(self.app, host=host, port=port, debug=debug)

if __name__ == '__main__':
    server = TerminalServer()
    server.run()